import gradio as gr
import numpy as np
from PIL import Image
import joblib
import sys
import subprocess
from pathlib import Path

# Ensure model path is absolute (same directory as this script)
BASE_DIR = Path(__file__).parent
MODEL_PATH = BASE_DIR / 'best_knn_model.pkl'

def load_or_train_model():
    """Try to load the pretrained model. If missing, run optimal_knn.py to train and save it,
    then load again. Uses the same Python interpreter (sys.executable).
    """
    try:
        return joblib.load(MODEL_PATH)
    except FileNotFoundError:
        print(f"Model file not found at {MODEL_PATH}. Attempting to train by running optimal_knn.py...")
        training_script = BASE_DIR / 'optimal_knn.py'
        if not training_script.exists():
            raise FileNotFoundError(f"Training script not found: {training_script}. Please create the model manually.")
        # Run the training script with the same interpreter
        completed = subprocess.run([sys.executable, str(training_script)], cwd=str(BASE_DIR))
        if completed.returncode != 0:
            raise RuntimeError(f"Training script failed with return code {completed.returncode}")
        # try loading again
        return joblib.load(MODEL_PATH)

# Load (or train if missing) model
model = load_or_train_model()

# 定义预测函数
def predict(my_dict):
    a = my_dict['composite']
    a = np.array(a)
    # 取alpha通道，转灰度
    a = a[:, :, 3]
    # 反色（黑底白字转为白底黑字，MNIST风格）
    a = 255 - a
    # 转为PIL灰度图
    b = Image.fromarray(a, 'L')
    # resize为8x8，与训练时一致
    c = b.resize((8, 8))
    # 像素缩放到0~16
    arr = np.array(c) / 255.0 * 16
    # 二值化处理，让输入更接近load_digits风格
    arr = (arr > 8) * 16
    # 可视化输入，便于调试
    Image.fromarray(arr.astype(np.uint8), 'L').save('input_debug.png')
    # 拉平成一维
    raveled_vector = arr.ravel()
    # 使用加载的模型进行预测
    prediction = model.predict([raveled_vector])
    # 返回预测结果
    return int(prediction[0])

# 创建 Gradio 界面
iface = gr.Interface(
    fn=predict,  # 设置要调用的预测函数
    inputs=gr.Sketchpad(),  # 创建手写板
    outputs=gr.Label(num_top_classes=1),
    title="knn手写数字识别",  # 设置界面的标题
    description="knn预测手写数字"  # 设置界面的描述
)

# 尝试创建共享链接；如果 frpc 缺失或下载被拦截（导致无法创建共享链接），
# 回退到本地-only 模式避免程序崩溃。
try:
    iface.launch(share=True)
except Exception as e:
    # 打印错误以便用户诊断（通常是 frpc 下载/权限/网络问题）
    print("Could not create share link. Falling back to local-only mode. Error:", e)
    # 本地模式
    iface.launch(share=False)